package engine

import (
	"sync"
	"sync/atomic"
)

// LoadBalancer 负载均衡器
type LoadBalancer struct {
	backends []*Backend
	current  uint64
	mu       sync.RWMutex
}

func NewLoadBalancer(cfg *BalancerConfig) (*LoadBalancer, error) {
	lb := &LoadBalancer{}
	for _, url := range cfg.Backends {
		backend := &Backend{
			URL: url,
		}
		backend.Health.Store(true)
		lb.backends = append(lb.backends, backend)
	}
	return lb, nil
}

// Next 获取下一个后端服务器
func (lb *LoadBalancer) Next() *Backend {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	if len(lb.backends) == 0 {
		return nil
	}

	next := atomic.AddUint64(&lb.current, 1)
	idx := next % uint64(len(lb.backends))
	return lb.backends[idx]
}

// MarkFailed 标记后端服务器失败
func (lb *LoadBalancer) MarkFailed(url string) {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	for _, b := range lb.backends {
		if b.URL == url {
			b.Health.Store(false)
			b.FailCount.Add(1)
			break
		}
	}
}

// MarkAlive 标记后端服务器恢复
func (lb *LoadBalancer) MarkAlive(url string) {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	for _, b := range lb.backends {
		if b.URL == url {
			b.Health.Store(true)
			b.FailCount.Store(0)
			break
		}
	}
}

// UpdateBackends 更新后端服务器列表
func (lb *LoadBalancer) UpdateBackends(urls []string) {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	newBackends := make([]*Backend, 0, len(urls))
	for _, url := range urls {
		// 检查是否已存在相同的后端
		var existing *Backend
		for _, b := range lb.backends {
			if b.URL == url {
				existing = b
				break
			}
		}

		if existing != nil {
			newBackends = append(newBackends, existing)
		} else {
			backend := &Backend{URL: url}
			backend.Health.Store(true)
			newBackends = append(newBackends, backend)
		}
	}

	lb.backends = newBackends
}

// Close 关闭负载均衡器并释放资源
func (lb *LoadBalancer) Close() {
	lb.mu.Lock()
	defer lb.mu.Unlock()
	lb.backends = nil
}
